package cn.iocoder.yudao.framework.websocket.core.handler;

import cn.iocoder.yudao.framework.websocket.config.WebSocketProperties;
import cn.iocoder.yudao.framework.websocket.core.SemaphoreUtils;
import cn.iocoder.yudao.framework.websocket.core.WebSocketSessionHandler;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.SubscribableChannel;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;

import java.util.concurrent.Semaphore;

@Slf4j
public class CustomSubProtocolWebSocketHandler extends SubProtocolWebSocketHandler {

    private final WebSocketProperties webSocketProperties;

    private final Semaphore socketSemaphore ;
    public CustomSubProtocolWebSocketHandler(MessageChannel clientInboundChannel, SubscribableChannel clientOutboundChannel, WebSocketProperties webSocketProperties) {
        super(clientInboundChannel, clientOutboundChannel);
        this.webSocketProperties = webSocketProperties;
        this.socketSemaphore = new Semaphore(webSocketProperties.getMaxConnection());
    }


    @Override
    public void afterConnectionEstablished( WebSocketSession session) throws Exception {
        super.afterConnectionEstablished(session);
        log.info(session.getAcceptedProtocol());
        // 尝试获取信号量
        boolean semaphoreFlag = SemaphoreUtils.tryAcquire(socketSemaphore);
        if (!semaphoreFlag)
        {
            // 未获取到信号量
            log.error("\n 当前在线人数超过限制数- {}", webSocketProperties.getMaxConnection());
            WebSocketSessionHandler.sendMessageToUserByText(session, "当前在线人数超过限制数：" + webSocketProperties.getMaxConnection());
            session.close();
        }
        else
        {
            // 添加用户
            WebSocketSessionHandler.put(session.getId(), session);
            log.info("\n 建立连接 - {}", session);
            log.info("\n 当前人数 - {}", WebSocketSessionHandler.getUsers().size());
          //  WebSocketSessionHandler.sendMessageToUserByText(session, "连接成功");
        }
    }
    @Override
    public void afterConnectionClosed( WebSocketSession session, CloseStatus closeStatus) throws Exception {
        super.afterConnectionClosed(session,closeStatus);
        log.info("\n 关闭连接 - {}", session);
        // 移除用户
        WebSocketSessionHandler.remove(session.getId());
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }
    @Override
    public void handleTransportError(WebSocketSession session,  Throwable exception) throws Exception {
        super.handleTransportError(session,exception);
        if (session.isOpen())
        {
            // 关闭连接
            session.close();
        }
        String sessionId = session.getId();
        log.info("\n 连接异常 - {}", sessionId);
        log.info("\n 异常信息 - {}", exception.getMessage());
        // 移出用户
        WebSocketSessionHandler.remove(sessionId);
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }

    @Override
    public void handleMessage(WebSocketSession session,  WebSocketMessage<?> message) throws Exception {
       super.handleMessage(session ,message);
       //TODO 这里可以预处理websocket消息
    }
}
